from langchain.vectorstores import Chroma
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
import torch, gc, os
from utils.QwenLM import QwenLM
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler  # for streaming response
from langchain.callbacks.manager import CallbackManager

callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
# from utils.prompt_template import get_prompt_template
from config.index import (
    EMBEDDING_MODEL_NAME,
    Persist_dir,
    MODEL_NAME,
    EMBEDDING_DEVICE
)


# 定义 Embeddings
embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME, model_kwargs={"device": EMBEDDING_DEVICE})

# 加载数据库
vectordb = Chroma(
    persist_directory=Persist_dir, 
    embedding_function=embeddings
)


llm = QwenLM(model_path = MODEL_NAME)
# 我们所构造的 Prompt 模板
template = """你是一个sql大神。如果你不知道答案，就说你不知道，不要试图编造答案。总是在回答的最后说“谢谢你的提问！”。
Context: {history} 
（如果不相关，您无需使用这些信息）
\n {context}
问题: {question}
回答:"""
# memory 历史记录
memory = ConversationBufferMemory(input_key="question", memory_key="history")
# 调用 LangChain 的方法来实例化一个 Template 对象，该对象包含了 context 和 question 两个变量，在实际调用时，这两个变量会被检索到的文档片段和用户提问填充
prompt = PromptTemplate(input_variables=["history", "context", "question"],template=template)
qa_chain = RetrievalQA.from_chain_type(
    llm, 
    chain_type="stuff",
    retriever=vectordb.as_retriever(),
    return_source_documents=True,
    callbacks=callback_manager,
    chain_type_kwargs={"prompt": prompt, "memory": memory},
)

while True:
    query = input("\n你说: ")
    if query == "exit":
        break
    result  = qa_chain({"query": query, })
    print("机器人：", result)

    # 用LLM 直接回答
    # result_2 = llm(question)
    # print("大模型回答 的结果：")
    # print(result_2)